
from pydantic import Field
from core.llm.base_analyze import BaseAnalyzer
from core import prompt_dir
from typing import Dict
from loguru import logger


class UserInfoAnalyzer(BaseAnalyzer):

    def __init__(self):
        super().__init__(prompt_file=prompt_dir / '用户信息抽取.md')
        self.keys = ['org', 'leadOrgs', 'primaryCategory', 'category', 'positions', 'roles']

    def invoke(
            self,
            user_data: Dict = Field(description="用户信息"),
            query: str = Field(description="用户问题描述")
    ):
        prompt = self.template_prompt.format(
            inputs={
                "user_data": user_data
            },
            remove_template_variables=True
        )
        data = self.analyze(prompt)
        return data
